#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <string.h>
#include <iostream>

bool generateKeyToFile() {
    int result = 0;
    RSA *rsa = NULL;
    BIGNUM *bne = NULL;
    BIO *bp_public = NULL, *bp_private = NULL;

    RSA *publicKey = NULL;
    RSA *privateKey = NULL;

    int bits = 1024;
    unsigned long e = RSA_F4;

    // 1. generate rsa key
    bne = BN_new();
    result = BN_set_word(bne, e);
    if (result != 1) {
        goto free_all;
    }

    rsa = RSA_new();
    result = RSA_generate_key_ex(rsa, bits, bne, NULL);
    if (result != 1) {
        goto free_all;
    }

    publicKey = RSAPublicKey_dup(rsa);
    privateKey = RSAPrivateKey_dup(rsa);
    RSA_print_fp(stdout, rsa, 5);
    printf("\n\n");
    RSA_print_fp(stdout, publicKey, 5);
    printf("\n\n");
    RSA_print_fp(stdout, privateKey, 5);
    printf("\n\n");
    RSA_free(publicKey);
    RSA_free(privateKey);

    // 2. save public key
    bp_public = BIO_new_file("public.pem", "w+");
    result = PEM_write_bio_RSAPublicKey(bp_public, rsa);
    if (result != 1) {
        goto free_all;
    }

    // 3. save private key
    bp_private = BIO_new_file("private.pem", "w+");
    result = PEM_write_bio_RSAPrivateKey(bp_private, rsa, NULL, NULL, 0, NULL, NULL);

    // 4. free
    free_all:

    BIO_free_all(bp_public);
    BIO_free_all(bp_private);
    RSA_free(rsa);
    BN_free(bne);

    return (result == 1);
}

RSA *generateKey() {
    int ret = 0;
    RSA *rsa = NULL;
    BIGNUM *bne = NULL;
    int bits = 1024;
    unsigned long e = RSA_F4;

    bne = BN_new();
    ret = BN_set_word(bne, e);
    if (ret != 1) {
        RSA_free(rsa);
        BN_free(bne);
    }

    rsa = RSA_new();
    ret = RSA_generate_key_ex(rsa, bits, bne, NULL);
    if (ret != 1) {
        RSA_free(rsa);
        BN_free(bne);
    }

    return rsa;
}

unsigned char *rsaEncrypt(unsigned char *originalText, RSA *rsa, int *len) {
    int rsa_len = RSA_size(rsa);
    unsigned char *cipherText = (unsigned char *) malloc(rsa_len + 1);
    memset(cipherText, 0, rsa_len + 1);
    int ret = RSA_public_encrypt(rsa_len, originalText, cipherText, rsa, RSA_NO_PADDING);
    if (ret < 0) {
        printf("Encrypt failed! ");
        return NULL;
    }
    *len = ret;
//    LOGW("rsaEncrypt ret : %d", ret);
//    LOGD("cipherText : %s\n", cipherText);
    return cipherText;
}

unsigned char *rsaDecrypt(unsigned char *cipherText, RSA *rsa, int *len) {
    int rsa_len = RSA_size(rsa);
    unsigned char *originalText = (unsigned char *) malloc(rsa_len + 1);
    memset(originalText, 0, rsa_len + 1);
    int ret = RSA_private_decrypt(rsa_len, cipherText, originalText, rsa, RSA_NO_PADDING);
    if (ret < 0) {
        printf("Decrypt failed! ");
        return NULL;
    }
    *len = ret;
//    LOGW("rsaDecrypt ret : %d", ret);
//    LOGD("originalText : %s\n", originalText);
    return originalText;
}

char *readTxt(const char *filePath) {
    FILE *file = fopen(filePath, "r");
    fseek(file, 0, SEEK_END);
    long fileLen = ftell(file);
    fseek(file, 0, SEEK_SET);
    char *text = (char *) malloc(fileLen + 1);
    memset(text, 0, fileLen + 1);
    fread(text, sizeof(char), fileLen, file);
    fclose(file);
    return text;
}

RSA *charToPublicKey(char *strPublicKey) {
    BIO *bio = NULL;
    RSA *rsa = NULL;
    if ((bio = BIO_new_mem_buf(strPublicKey, -1)) == NULL) {
        printf("%s", "BIO_new_mem_buf failed!");
    }
    rsa = PEM_read_bio_RSAPublicKey(bio, NULL, NULL, NULL);
    if (!rsa) {
        ERR_load_crypto_strings();
        char errBuf[512];
        ERR_error_string_n(ERR_get_error(), errBuf, sizeof(errBuf));
        printf("%s%s%s", "load public key failed[", errBuf, "]");
    }
    BIO_free_all(bio);
    return rsa;
}

char *publicKeyToChar(RSA *publicKey) {
    BIO *bio = BIO_new(BIO_s_mem());
    int ret = PEM_write_bio_RSAPublicKey(bio, publicKey);
    if (ret != 1) {
        ERR_load_crypto_strings();
        char errBuf[512];
        ERR_error_string_n(ERR_get_error(), errBuf, sizeof(errBuf));
        printf("%s%s%s", "convert public key failed[", errBuf, "]");
    }
    int len = BIO_pending(bio);
    char *result = (char *) malloc(len + 1);
    memset(result, 0, len + 1);
    BIO_read(bio, result, len);
    BIO_free_all(bio);
    return result;
}

RSA *getPublicKeyFromFile(char *filePath) {
    FILE *fp;
    if ((fp = fopen(filePath, "r")) == NULL) {
        return NULL;
    }
    RSA *rsa = NULL;
    rsa = PEM_read_RSAPublicKey(fp, NULL, NULL, NULL);
    if (!rsa) {
        ERR_load_crypto_strings();
        char errBuf[512];
        ERR_error_string_n(ERR_get_error(), errBuf, sizeof(errBuf));
        printf("%s%s%s", "load public key failed[", errBuf, "]");
    }
    return rsa;
}

unsigned char *base64Encode(unsigned char *input, int length, bool with_new_line) {
    BIO *bmem = NULL;
    BIO *b64 = NULL;
    BUF_MEM *bptr = NULL;

    b64 = BIO_new(BIO_f_base64());
    if (!with_new_line) {
        BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
    }
    bmem = BIO_new(BIO_s_mem());
    b64 = BIO_push(b64, bmem);
    BIO_write(b64, input, length);
    BIO_flush(b64);
    BIO_get_mem_ptr(b64, &bptr);

    unsigned char *buff = (unsigned char *) malloc(bptr->length + 1);
    memcpy(buff, bptr->data, bptr->length);
    buff[bptr->length] = 0;

    BIO_free_all(b64);

    return buff;
}

unsigned char *base64Decode(unsigned char *input, int length, bool with_new_line) {
    BIO *b64 = NULL;
    BIO *bmem = NULL;
    unsigned char *buffer = (unsigned char *) malloc(length);
    memset(buffer, 0, length);

    b64 = BIO_new(BIO_f_base64());
    if (!with_new_line) {
        BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
    }
    bmem = BIO_new_mem_buf(input, length);
    bmem = BIO_push(b64, bmem);
    BIO_read(bmem, buffer, length);

    BIO_free_all(bmem);

    return buffer;
}

int main(int argc, char *argv[]) {
    RSA *rsa = generateKey();

    RSA *publicKey = NULL;
    RSA *privateKey = NULL;

    publicKey = RSAPublicKey_dup(rsa);
    privateKey = RSAPrivateKey_dup(rsa);
//    RSA_print_fp(stdout, rsa, 5);
//    printf("\n\n");
//    RSA_print_fp(stdout, publicKey, 5);
//    printf("\n\n");
//    RSA_print_fp(stdout, privateKey, 5);
//    printf("\n\n");

    unsigned char *originalText = (unsigned char *) "abcdefg";

    printf("OriginalText :%s\n", originalText);

    BIO *bp_public = NULL;
    bp_public = BIO_new_file("public.pem", "w+");
    PEM_write_bio_RSAPublicKey(bp_public, publicKey);
    BIO_free_all(bp_public);

    publicKey = charToPublicKey(readTxt("public.pem"));

    char *charText = publicKeyToChar(publicKey);
    printf("charText :\n%s", charText);

    publicKey = charToPublicKey(charText);

    int len;
    unsigned char *cipherText = rsaEncrypt(originalText, publicKey, &len);
    int ret = (int) strlen((const char *) cipherText);
    printf("Size:%d\n", ret);
    printf("CipherText:");
    for (int i = 0; i < ret; i++) {
        printf("%02x", cipherText[i]);
    }
    printf("\n");
    printf("\n");

    unsigned char *base64CipherText = base64Encode(cipherText, len, false);
    printf("base64Text :\n%s", base64CipherText);
    printf("\n");
    printf("\n");

    unsigned char *cipherTextFromBase64 = base64Decode(base64CipherText, strlen((const char *) base64CipherText),
                                                       false);
    printf("cipherTextFromBase64 :\n%s", cipherTextFromBase64);
    printf("\n");
    printf("\n");

    unsigned char *decryptText = rsaDecrypt(cipherTextFromBase64, privateKey, &len);
    ret = (int) strlen((const char *) decryptText);
    printf("Size:%d\n", ret);
    printf("OriginalText :%s\n", decryptText);

    free(cipherText);
    free(base64CipherText);
    free(cipherTextFromBase64);
    free(decryptText);
    RSA_free(publicKey);
    RSA_free(privateKey);
    RSA_free(rsa);

    return 0;
}